[ExecuTorch][WebGPU] Add fp32 matmul ops (mm, bmm, linear) to the WebGPU backend#20917
[ExecuTorch][WebGPU] Add fp32 matmul ops (mm, bmm, linear) to the WebGPU backend#20917JCNTH wants to merge 2 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20917
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
Stack from ghstack (oldest at bottom):
Add fp32 GEMM handlers (
aten.mm,aten.bmm,aten.linear) — the dense matmuls the on-device training tail needs (tiled + vec4 WGSL).Key changes:
runtime/ops/{mm,bmm,linear}/— tiled + vec4 fp32 GEMM WGSL kernels + handlersCMakeLists.txtWEBGPU_SRCS— wire the three sourcesReuses the shared Vulkan partitioner (
aten.mm/bmm/linearalready have VulkanOpFeatures); this adds the WebGPU kernels only.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755135
Differential Revision: D111755135